// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#include "utils/model_config.h"

#include "base_util/utils.h"
#include "postprocess_op.h"
#include "preprocess_op.h"
#include "ocr_det.h"

namespace paddle_infer {


OCRDetector::OCRDetector(LogInfo *log_info):
  log_ifo(log_info) {   
  // resize_op_ = new ResizeImgType0();
  // normalize_op_ = new Normalize0();
  // permute_op_ = new Permute();

  // post-process
  post_processor_ = new OCRPostProcessor();
}

int OCRDetector::Run(cv::Mat& img, 
                      std::vector<BaseInfoPD*> &text_lines, 
                      PDModelManager *mng,
                      int model_idx )
{

  if (mng->configs[model_idx]->algorithm == pd_text_det || mng->configs[model_idx]->algorithm == pd_plate_rec){
    return RunDB(img, text_lines, mng, model_idx);
  }else{
    return pd_not_support_algorithm;
  }
  return pd_not_support_algorithm;
}

int OCRDetector::RunDB(cv::Mat& img, 
                        std::vector<BaseInfoPD*> &rec_ifos, 
                        PDModelManager *mng,
                        int md_idx ) 
{
  int ret=0;
  std::vector<TextDetInfoPD> text_lines;
  if (img.empty()){return PD_IMAGE_EMPTY_ERROR;}
  auto& cfg = mng->configs[md_idx];
  ImageBlob img_blob = ImageBlob();
  cfg->transforms->run(img, img_blob, cfg);
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.1 RunDB");

  auto input_names = mng->predictors[md_idx]->GetInputNames();
  auto in_tensor = mng->predictors[md_idx]->GetInputTensor(input_names[0]);
  in_tensor->Reshape({1, cfg->channels, img_blob.new_im_shape[0], img_blob.new_im_shape[1]});
  in_tensor->CopyFromCpu(img_blob.im_data.data());

  mng->predictors[md_idx]->ZeroCopyRun();
  if (log_ifo->log_level_4) spdlog::get("logger")->info("1.2.2.2 RunDB");

  auto output_names = mng->predictors[md_idx]->GetOutputNames();
  if (log_ifo->log_level_4) { log_tensor_shape("output_names", output_names); }
  
  auto out_tensor = mng->predictors[md_idx]->GetOutputTensor(output_names[0]);
  std::vector<int> output_shape = out_tensor->shape();
  if (log_ifo->log_level_4) { log_tensor_shape("output_tensor", output_shape); }
 
  int size = std::accumulate(output_shape.begin(), output_shape.end(), 1, std::multiplies<int>());
  std::vector<float> output_data(size);
  out_tensor->CopyToCpu(output_data.data());

  int n2 = output_shape[2];
  int n3 = output_shape[3];
  int n = n2 * n3;

  std::vector<float> pred(n, 0.0);
  std::vector<unsigned char> cbuf(n, ' ');

  for (int i = 0; i < n; i++) {
    pred[i] = float(output_data[i]);
    cbuf[i] = (unsigned char)((output_data[i]) * 255);
  }

  cv::Mat cbuf_map(n2, n3, CV_8UC1, (unsigned char *)cbuf.data());
  cv::Mat pred_map(n2, n3, CV_32F, (float *)pred.data());

  const double threshold = mng->configs[md_idx]->det_db_thresh * 255;
  const double maxvalue = 255;
  cv::Mat bit_map;
  cv::threshold(cbuf_map, bit_map, threshold, maxvalue, cv::THRESH_BINARY);
  if (log_ifo->log_level_4){
    std::string save_path = bf::path(log_ifo->get_log_path()).append(std::to_string(log_ifo->counter()) + "_bit_map.jpg").string();
    cv::imwrite(save_path, bit_map);
  }

  ret = post_processor_->BoxesFromBitmap( text_lines, pred_map, bit_map, 
    mng->configs[md_idx]->det_db_box_thresh, 
    mng->configs[md_idx]->det_db_unclip_ratio,
    mng->configs[md_idx]->det_expand_by_min_score);

  ret = post_processor_->FilterTagDetRes(text_lines, img_blob.scale, img_blob.scale, img);

  for (auto& tl : text_lines) { rec_ifos.push_back(new TextDetInfoPD(tl.four_points, tl.score)); }
  return ret;
}


// int OCRDetector::RunDB(cv::Mat& img, 
//                         std::vector<BaseInfoPD*> &rec_ifos, 
//                         PDModelManager *mng,
//                         int model_idx ) 
// {
//   int ret=0;
//   float ratio_h{};
//   float ratio_w{};
//   std::vector<TextDetInfoPD> text_lines;

//   cv::Mat srcimg;
//   cv::Mat resize_img;
//   img->copyTo(srcimg);
//   this->resize_op_->Run(*img, resize_img, mng->configs[model_idx]->max_side_len, 
//                         ratio_h, ratio_w);

//   this->normalize_op_->Run(&resize_img, this->mean_, this->scale_,
//                           this->is_scale_);

//   std::vector<float> input(1 * 3 * resize_img.rows * resize_img.cols, 0.0f);
//   this->permute_op_->Run(&resize_img, input.data());

//   auto input_names = mng->predictors[model_idx]->GetInputNames();
//   auto input_t = mng->predictors[model_idx]->GetInputTensor(input_names[0]);
//   input_t->Reshape({1, 3, resize_img.rows, resize_img.cols});
//   input_t->CopyFromCpu(input.data());

//   mng->predictors[model_idx]->ZeroCopyRun();

//   std::vector<float> out_data;
//   auto output_names = mng->predictors[model_idx]->GetOutputNames();
//   auto output_t = mng->predictors[model_idx]->GetOutputTensor(output_names[0]);
//   std::vector<int> output_shape = output_t->shape();
//   int out_num = std::accumulate(output_shape.begin(), output_shape.end(), 1,
//                                 std::multiplies<int>());

//   out_data.resize(out_num);
//   output_t->CopyToCpu(out_data.data());

//   int n2 = output_shape[2];
//   int n3 = output_shape[3];
//   int n = n2 * n3;

//   std::vector<float> pred(n, 0.0);
//   std::vector<unsigned char> cbuf(n, ' ');

//   for (int i = 0; i < n; i++) {
//     pred[i] = float(out_data[i]);
//     cbuf[i] = (unsigned char)((out_data[i]) * 255);
//   }

//   cv::Mat cbuf_map(n2, n3, CV_8UC1, (unsigned char *)cbuf.data());
//   cv::Mat pred_map(n2, n3, CV_32F, (float *)pred.data());

//   const double threshold = mng->configs[model_idx]->det_db_thresh * 255;
//   const double maxvalue = 255;
//   cv::Mat bit_map;
//   cv::threshold(cbuf_map, bit_map, threshold, maxvalue, cv::THRESH_BINARY);
//   if (log_info_->log_level_3){
//     cv::imwrite(log_info_->get_log_path()+"/" \
//                               +std::to_string(log_info_->counter()) \
//                               +"_bit_map.jpg", bit_map);
//   }

//   ret = post_processor_->BoxesFromBitmap( text_lines, pred_map, bit_map, 
//     mng->configs[model_idx]->det_db_box_thresh, 
//     mng->configs[model_idx]->det_db_unclip_ratio,
//     mng->configs[model_idx]->det_expand_by_min_score);

//   ret = post_processor_->FilterTagDetRes(text_lines, ratio_h, ratio_w, srcimg);

//   for (auto& tl : text_lines) { rec_ifos.push_back((BaseInfoPD*)&tl); }
//   return ret;
// }

} // namespace paddle_infer
